import os
import cv2
import numpy as np
from sklearn.utils import shuffle
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import models, layers

TRAFFIC_LIGHT_RED      = 0
TRAFFIC_LIGHT_GREEN    = 1
TRAFFIC_LIGHT_VEHICLE  = 2

IMAGE_HEIGHT  = 32
IMAGE_WIDTH   = 32
IMAGE_CHANNEL = 3

saved_model_name = 'guide_cnn.h5'

def get_light_name(label):
    if label == TRAFFIC_LIGHT_RED:
        return "red"

    if label == TRAFFIC_LIGHT_GREEN:
        return "green"

    if label == TRAFFIC_LIGHT_VEHICLE:
        return "vehicle"

def get_model():
    model = models.Sequential()
    
    model.add(layers.Conv2D(32, (5, 5), 
                input_shape=(IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNEL)))
    model.add(layers.LeakyReLU(alpha=0.5))
    model.add(layers.MaxPooling2D((2, 2)))
    
    model.add(layers.Conv2D(48, (3, 3)))
    model.add(layers.LeakyReLU(alpha=0.5))
    model.add(layers.MaxPooling2D((2, 2)))
    
    model.add(layers.Conv2D(64, (3, 3)))
    model.add(layers.LeakyReLU(alpha=0.5))
    model.add(layers.MaxPooling2D((2, 2)))
    
    model.add(layers.Flatten())

    model.add(layers.Dense(100))
    model.add(layers.LeakyReLU(alpha=0.5))
    
    model.add(layers.Dense(50))
    model.add(layers.LeakyReLU(alpha=0.5))
    
    model.add(layers.Dense(3))
    model.add(layers.LeakyReLU(alpha=0.5))
    
    model.summary()

    return model

def read_image(path):
    image = cv2.imread(path)
    image = cv2.resize(image, (IMAGE_WIDTH, IMAGE_HEIGHT))
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    normalizer_image = image / 255.0 - 0.5

    return normalizer_image

def get_signal_image_data():

    # train
    train_images = []
    train_labels = []
    train_image_names = []

    train_path = './traffic_light_images/train'
    
    path = train_path + '/red'
    for file_name in os.listdir( path ):
        path_name = path + '/' + file_name
        train_images.append( read_image(path_name) )
        train_labels.append( TRAFFIC_LIGHT_RED )
        train_image_names.append( file_name )
    
    path = train_path + '/green'
    for file_name in os.listdir( path ):
        path_name = path + '/' + file_name
        train_images.append( read_image(path_name) )
        train_labels.append( TRAFFIC_LIGHT_GREEN )
        train_image_names.append( file_name )
    
    path = train_path + '/vehicle'
    for file_name in os.listdir( path ):
        path_name = path + '/' + file_name
        train_images.append( read_image(path_name) )
        train_labels.append( TRAFFIC_LIGHT_VEHICLE )
        train_image_names.append( file_name )



    return shuffle(train_images, train_labels, train_image_names)

def train():
    model = get_model()

    model.compile(optimizer='adam',
                  loss=losses.SparseCategoricalCrossentropy(from_logits=True),
                  metrics=['accuracy']) 

    (train_images, train_labels, _) = get_signal_image_data()

    model.fit( np.array(train_images), np.array(train_labels), epochs=1 )

    #models.save_model( model=model, filepath=saved_model_name, save_format='tf' )
    models.save_model( model, saved_model_name )


def load_model():
    model = models.load_model( saved_model_name ) 
    model = models.Sequential( [model, layers.Softmax()])
    return model

def predict(model, image):
    image = cv2.resize(image, (IMAGE_WIDTH, IMAGE_HEIGHT))
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    normalizer_image = image / 255.0 - 0.5

    input = normalizer_image.reshape(1, IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNEL)
    pred_results = model.predict( input )
    signal_light = np.argmax( pred_results[0] )
    return signal_light

if __name__ == "__main__":

    train()

    model = models.load_model( saved_model_name ) 
    #model = get_model()
    #model.load_weights( 'my_model' )
    model = Sequential( [model, layers.Softmax()])

    # test
    test_path = './traffic_light_images/test'
    total_count = 0
    correct_count = 0
    path = test_path + '/red'
    for file_name in os.listdir( path ):
        path_name = path + '/' + file_name
        image = cv2.imread( path_name )
        if image is None:
            print( "read image %s failed" % (path_name) )
            break
        signal_light = predict( model, image )
        total_count = total_count + 1
        if signal_light == TRAFFIC_LIGHT_RED:
            correct_count = correct_count + 1
        print( "%s:  is %s, result is %s" % (file_name, get_light_name(TRAFFIC_LIGHT_RED), get_light_name(signal_light)))
    
    path = test_path + '/green'
    for file_name in os.listdir( path ):
        path_name = path + '/' + file_name
        image = cv2.imread( path_name )
        if image is None:
            print( "read image failed" )
            break
        signal_light = predict( model, image )
        total_count = total_count + 1
        if signal_light == TRAFFIC_LIGHT_GREEN:
            correct_count = correct_count + 1
        print( "%s:  is %s, result is %s" % (file_name, get_light_name(TRAFFIC_LIGHT_GREEN), get_light_name(signal_light)))
    
    path = test_path + '/vehicle'
    for file_name in os.listdir( path ):
        path_name = path + '/' + file_name
        image = cv2.imread( path_name )
        if image is None:
            print( "read image failed" )
            break
        signal_light = predict( model, image )
        total_count = total_count + 1
        if signal_light == TRAFFIC_LIGHT_VEHICLE:
            correct_count = correct_count + 1
        print( "%s:  is %s, result is %s" % (file_name, get_light_name(TRAFFIC_LIGHT_VEHICLE), get_light_name(signal_light)))

    print( "total = %d, correct_count = %d" % (total_count, correct_count) )
    print( "accuracy is %f" % (float(correct_count) / float(total_count)) )
